import uuid
import json
from redis.asyncio import Redis
from config.env import RedisConfig
from ocr_core.schema import LlmTask, LlmTaskResult, TaskStatus
from utils.log_util import logger
from config.mongodb import MongoDB, LlmTaskMDB, LlmResultMDB


class LlmProducer:
    """LLM 任务生产者，负责推送任务到 Redis 队列并查询任务结果"""

    def __init__(self, redis_client: Redis):
        self.redis: Redis = redis_client

    async def create_llm_task(self, task: LlmTask) -> str:
        """
        创建新任务并推送到 Redis 队列
        Args:
            task: LlmTask 任务数据
        Returns:
            task_id: 生成或传入的任务ID
        """
        task.task_id = task.task_id or str(uuid.uuid4())

        try:
            serialized_task = json.dumps(task.model_dump(), ensure_ascii=False)
        except Exception as e:
            logger.error(f"任务序列化失败: {task}, 错误: {e}")
            raise

        # 初始化任务状态
        init_result = LlmTaskResult(
            task_id=task.task_id,
            scenario=task.scenario,
            status=TaskStatus.PENDING,
            error="",
            data="",
            llm_ms=0,
        )

        success = await self.redis.set(
            f"{RedisConfig.redis_llm_result}{task.task_id}",
            json.dumps(init_result.model_dump(), ensure_ascii=False),
            ex=3600,  # 1小时过期
        )

        if not success:
            logger.error(f"任务 {task.task_id} 状态初始化失败")
            raise RuntimeError("任务状态初始化失败")

        logger.debug(f"任务 {task.task_id} 创建成功 ✅")

        # 推送任务到队列（默认 FIFO）
        await self.redis.rpush(RedisConfig.redis_llm_queue, serialized_task)

        # 插入mongodb的llm文档
        LlmResultMDB.insert_one(init_result.model_dump())
        LlmTaskMDB.insert_one(task.model_dump())
        return task.task_id

    async def get_llm_task_result(self, task_id: str) -> LlmTaskResult:
        """
        获取任务状态
        Args:
            task_id: 任务ID
        Returns:
            LlmTaskResult: 包含任务状态和结果
        """
        find_one = LlmResultMDB.find_one({"task_id": task_id})
        if find_one:
            try:
                return LlmTaskResult(**find_one)
            except Exception as e:
                logger.error(f"任务结果反序列化失败 [task_id={task_id}] 错误: {e}")

        # 兜底返回
        return LlmTaskResult(
            task_id=task_id,
            scenario="",
            status=TaskStatus.NOT_FOUND,
            data="",
            llm_ms=0,
            error="Task not found",
        )
